package chapter11

import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.SparkSession

import scala.collection.mutable.WrappedArray

/**
 * author: yuhui
 * descriptions:
 * 1）通过一个人的age,height,weight,yanzhi,score，作为特征
 * 2）使用余弦相似度函数计算出相似度
 * date: 2024 - 11 - 28 3:45 下午
 */
object UDF03 {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    import spark.implicits._

    // 加载数据
    val df = spark.read.option("inferSchema", true).option("header", true).csv("BookData/input/11features.csv")
    df.show()

    // 将数据转换为包含特征数组的DataFrame
    val featuresDF = df.selectExpr("id", "name", "array(age,height,weight,yanzhi,score) as features")

    // 将表自己和自己join，得到每个人和其他所有人的连接行
    val joined = featuresDF.as("a").join(featuresDF.as("b"), $"a.id" < $"b.id")

    // 定义一个计算余弦相似度的函数
    val cosinSim: UserDefinedFunction = udf((f1: WrappedArray[Double], f2: WrappedArray[Double]) => {
      val norm1 = Math.sqrt(f1.map(Math.pow(_, 2)).sum)
      val norm2 = Math.sqrt(f2.map(Math.pow(_, 2)).sum)
      val dotProduct = f1.zip(f2).map(p => p._1 * p._2).sum
      dotProduct / (norm1 * norm2)
    })

    // 注册UDF
    spark.udf.register("cos_sim", cosinSim)

    // 创建临时视图并计算余弦相似度
    joined.select($"a.id".as("aid"), $"a.features".as("afeatures"),
      $"b.id".as("bid"), $"b.features".as("bfeatures"))
      .createTempView("joined_temp")

    // 使用UDF直接在DataFrame SQL中计算余弦相似度
    spark.sql("SELECT aid, bid, cos_sim(afeatures, bfeatures) AS cos_sim_result FROM joined_temp").show()

    // 使用UDF直接在DataFrame API中计算余弦相似度
    joined.select($"a.id".as("aid"), $"b.id".as("bid"), cosinSim($"a.features", $"b.features").as("cos_sim_result")).show()

    // 关闭Spark
    spark.close()
  }
}